#!/usr/bin/env python3

import rclpy
from rclpy.node import Node
from rclpy.executors import ExternalShutdownException
from rclpy.qos import qos_profile_sensor_data
from std_msgs.msg import String, Bool
from geometry_msgs.msg import Pose, PoseStamped
from rclpy.action import ActionClient
from moveit_msgs.action import MoveGroup
from moveit_msgs.msg import MoveItErrorCodes
from tf2_ros import Buffer, TransformListener
import numpy as np
import copy
import time

# 全局变量
left_pose = None  # 接收到的纸巾目标位姿
yogurt_pose = None  # 接收到的酸奶目标位姿
last_valid_pose = None
last_valid_pose_yogurt = None

start_move = False  # 是否开始任务
task_state = "stand_by"  # 任务状态：stand_by, move_to_paper, grasp, move_to_release, release, return_home_after_paper, move_to_yogurt, grasp_yogurt, move_to_yogurt_release, release_yogurt, return_home_after_yogurt

# MoveIt 组
move_group_client = None

# 手控制发布器
hand_control_publisher = None

# 纸巾抓取备选位姿
paper_alternative_poses = [
    {
        "orientation": {
            "x": -0.0321,
            "y": 0.1140,
            "z": 0.9236,
            "w": 0.3645
        }
    },
    {
        "orientation": {
            "x": -0.014,
            "y": 0.332,
            "z": 0.912,
            "w": 0.24
        }
    }
]

# 酸奶抓取备选位姿
yogurt_alternative_poses = [
    {
        "orientation": {
            "x": 0.6855,
            "y": -0.1520,
            "z": 0.6540,
            "w": 0.2812
        }
    },
    {
        "orientation": {
            "x": 0.6402,
            "y": -0.2407,
            "z": 0.6268,
            "w": 0.3731
        }
    }
]

# 固定的释放位置
release_pose = PoseStamped()
release_pose.header.frame_id = "base_link"
release_pose.pose.position.x = 0.5  # 固定释放位置的 x 坐标
release_pose.pose.position.y = -0.55  # 固定释放位置的 y 坐标
release_pose.pose.position.z = 0.85  # 固定释放位置的 z 坐标
release_pose.pose.orientation.x = 0.11236370921869202  # 固定释放位置的姿态
release_pose.pose.orientation.y = 0.15079563957693937
release_pose.pose.orientation.z = 0.7905900885524749
release_pose.pose.orientation.w = 0.5827541366801177


class TaskManager(Node):
    def __init__(self):
        super().__init__('task_manager_node')
        self.tf_buffer = Buffer()
        self.tf_listener = TransformListener(self.tf_buffer)

        self.create_subscription(PoseStamped, '/paper_pose', self.left_pose_callback, qos_profile_sensor_data)
        self.create_subscription(PoseStamped, '/yoger_pose', self.yogurt_pose_callback, qos_profile_sensor_data)
        self.create_subscription(Bool, '/start_move', self.start_move_callback, qos_profile_sensor_data)

        self.hand_control_publisher = self.create_publisher(String, "/hand_control_topic", qos_profile_sensor_data)

        self.move_group_client = ActionClient(self, MoveGroup, 'move_group')

    def left_pose_callback(self, msg):
        global left_pose, last_valid_pose
        if msg.pose.position.x == -1 or msg.pose.position.y == -1 or msg.pose.position.z == -1:
            if last_valid_pose:
                left_pose = copy.deepcopy(last_valid_pose)
        else:
            if last_valid_pose:
                current_position = np.array([msg.pose.position.x, msg.pose.position.y, msg.pose.position.z])
                last_position = np.array([last_valid_pose.pose.position.x, last_valid_pose.pose.position.y, last_valid_pose.pose.position.z])
                distance = np.linalg.norm(current_position - last_position)
                if distance > 0.10:
                    self.get_logger().info(f"Position change exceeds 1cm ({distance:.4f}m), skipping this pose.")
                    return
            left_pose = copy.deepcopy(msg)
            last_valid_pose = copy.deepcopy(msg)

    def yogurt_pose_callback(self, msg):
        """
        接收酸奶目标位姿的回调函数。
        """
        # 检查当前 pose 是否有效（避免接收到无效位姿）
        if msg.pose.position.x == -1 or msg.pose.position.y == -1 or msg.pose.position.z == -1:
            if self.last_valid_pose_yogurt:
                self.yogurt_pose = copy.deepcopy(self.last_valid_pose_yogurt)
            else:
                self.get_logger().warn('Invalid pose received, and no last valid pose exists.')
                return
        else:
            if self.last_valid_pose_yogurt:
                current_position = np.array([msg.pose.position.x, msg.pose.position.y, msg.pose.position.z])
                last_position = np.array(
                    [self.last_valid_pose_yogurt.pose.position.x, self.last_valid_pose_yogurt.pose.position.y, self.last_valid_pose_yogurt.pose.position.z])
                distance = np.linalg.norm(current_position - last_position)

                if distance > 0.10:
                    self.get_logger().info(f"Position change exceeds 1cm ({distance:.4f}m), skipping this pose.")
                    return
                else:
                    self.yogurt_pose = copy.deepcopy(msg)
                    self.last_valid_pose_yogurt = copy.deepcopy(msg)
            else:
                self.yogurt_pose = copy.deepcopy(msg)
                self.last_valid_pose_yogurt = copy.deepcopy(msg)

    def start_move_callback(self, msg):
        global start_move
        if msg.data:
            self.get_logger().info("Received start_move signal, starting task.")
            start_move = True
        else:
            self.get_logger().info("Received stop_move signal, stopping task.")
            start_move = False

    def move_to_pose(self, pose, link_name, is_yogurt=False):
        """
        移动手臂到指定位姿，使用分步规划策略。
        :param pose: Pose 消息，目标位姿
        :param link_name: 要移动的链接名称
        :param is_yogurt: 是否是酸奶任务
        :return: 是否成功
        """
        self.get_logger().info(f"Target pose: {pose}")
        current_pose = self.move_group.get_current_pose().pose
        self.get_logger().info(f"Current pose: {current_pose}")

        distance = np.linalg.norm(
            np.array([pose.position.x, pose.position.y, pose.position.z]) -
            np.array([current_pose.position.x, current_pose.position.y, current_pose.position.z])
        )
        self.get_logger().info(f"Distance to target pose: {distance}")

        step_size = 0.5
        if distance > step_size:
            num_steps = int(np.ceil(distance / step_size))

            for i in range(num_steps):
                interpolation_factor = (i + 1) / num_steps
                intermediate_pose = Pose()
                intermediate_pose.position.x = current_pose.position.x + (
                        pose.position.x - current_pose.position.x
                ) * interpolation_factor
                intermediate_pose.position.y = current_pose.position.y + (
                        pose.position.y - current_pose.position.y
                ) * interpolation_factor
                intermediate_pose.position.z = current_pose.position.z + (
                        pose.position.z - current_pose.position.z
                ) * interpolation_factor

                intermediate_pose.orientation = pose.orientation

                self.move_group.set_pose_target(intermediate_pose, link_name)
                plan = self.move_group.plan()

                if plan[0]:  # plan[0] 是是否成功生成计划
                    success = self.move_group.execute(plan[1], wait=True)
                    if not success:
                        self.get_logger().warn(f"Failed to move {link_name} to intermediate pose at step {i + 1}.")
                        if self.try_alternative_orientations(intermediate_pose, link_name, is_yogurt):
                            continue
                        else:
                            return False
                else:
                    self.get_logger().warn(f"Failed to generate plan for {link_name} at step {i + 1}.")
                    if self.try_alternative_orientations(intermediate_pose, link_name, is_yogurt):
                        continue
                    else:
                        return False

            self.get_logger().info(f"Successfully moved {link_name} to the target pose in {num_steps} steps.")
        else:
            self.move_group.set_pose_target(pose, link_name)
            plan = self.move_group.plan()

            if plan[0]:
                success = self.move_group.execute(plan[1], wait=True)
                if success:
                    self.get_logger().info(f"Successfully moved {link_name} to the target pose.")
                    return True
                else:
                    self.get_logger().warn(f"Failed to execute plan for {link_name}.")
                    return self.try_alternative_orientations(pose, link_name, is_yogurt)
            else:
                self.get_logger().warn(f"Failed to generate plan for {link_name}.")
                return self.try_alternative_orientations(pose, link_name, is_yogurt)

    def try_alternative_orientations(self, pose, link_name, is_yogurt=False):
        global paper_alternative_poses, yogurt_alternative_poses
        alternative_poses = yogurt_alternative_poses if is_yogurt else paper_alternative_poses
        for i, alt_pose in enumerate(alternative_poses):
            alternative_pose = copy.deepcopy(pose)
            alternative_pose.pose.orientation.x = alt_pose["orientation"]["x"]
            alternative_pose.pose.orientation.y = alt_pose["orientation"]["y"]
            alternative_pose.pose.orientation.z = alt_pose["orientation"]["z"]
            alternative_pose.pose.orientation.w = alt_pose["orientation"]["w"]
            if self.move_to_pose(alternative_pose, link_name, is_yogurt):
                self.get_logger().info(f"Successfully moved {link_name} to the target pose using alternative orientation {i + 1}.")
                return True
            else:
                self.get_logger().warn(f"Failed to move {link_name} to the target pose using alternative orientation {i + 1}.")
        self.get_logger().warn("Failed to move to target pose using all alternative orientations.")
        return False

    def move_to_stand_by_pose(self):
        stand_by_pose = PoseStamped()
        stand_by_pose.header.frame_id = "base_link"
        stand_by_pose.pose.position.x = 0.4
        stand_by_pose.pose.position.y = -0.55
        stand_by_pose.pose.position.z = 0.80
        stand_by_pose.pose.orientation.x = 0.11236370921869202
        stand_by_pose.pose.orientation.y = 0.15079563957693937
        stand_by_pose.pose.orientation.z = 0.7905900885524749
        stand_by_pose.pose.orientation.w = 0.5827541366801177
        return self.move_to_pose(stand_by_pose, "right_tool_link")

    def adjust_pose_for_paper(self, pose):
        adjusted_pose = copy.deepcopy(pose)
        adjusted_pose.pose.position.x += 0.02
        adjusted_pose.pose.position.y += -0.10
        adjusted_pose.pose.position.z += 0.03
        adjusted_pose.pose.orientation.x = 0.0927
        adjusted_pose.pose.orientation.y = 0.2553
        adjusted_pose.pose.orientation.z = 0.9106
        adjusted_pose.pose.orientation.w = 0.3113
        return adjusted_pose

    def adjust_pose_for_yogurt(self, pose):
        adjusted_pose = copy.deepcopy(pose)
        adjusted_pose.pose.position.x += 0.05
        adjusted_pose.pose.position.y += -0.05
        adjusted_pose.pose.position.z += -0.06
        adjusted_pose.pose.orientation.x = 0.6878
        adjusted_pose.pose.orientation.y = -0.0545
        adjusted_pose.pose.orientation.z = 0.6945
        adjusted_pose.pose.orientation.w = 0.2037
        return adjusted_pose

    def task_manager(self):
        global task_state, start_move, left_pose, yogurt_pose, hand_control_publisher, release_pose
        rate = self.create_rate(10)  # 10Hz
        while rclpy.ok():
            if not start_move:
                rate.sleep()
                continue

            if task_state == "stand_by":
                self.get_logger().info("Task state: stand_by")

                current_pose = move_group_client.get_current_pose().pose
                stand_by_pose = PoseStamped()
                stand_by_pose.header.frame_id = "base_link"
                stand_by_pose.pose.position.x = 0.4
                stand_by_pose.pose.position.y = -0.5
                stand_by_pose.pose.position.z = 0.80
                stand_by_pose.pose.orientation.x = 0.11236370921869202
                stand_by_pose.pose.orientation.y = 0.15079563957693937
                stand_by_pose.pose.orientation.z = 0.7905900885524749
                stand_by_pose.pose.orientation.w = 0.5827541366801177

                distance = np.linalg.norm(
                    np.array([current_pose.position.x, current_pose.position.y, current_pose.position.z]) -
                    np.array(
                        [stand_by_pose.pose.position.x, stand_by_pose.pose.position.y, stand_by_pose.pose.position.z])
                )

                # 判断是否需要移动
                if distance > 0.1:
                    self.get_logger().info("Arm is not in stand_by pose, moving to stand_by pose...")
                    if self.move_to_stand_by_pose():
                        self.get_logger().info("Arm is now in stand_by pose.")
                    else:
                        self.get_logger().warn("Failed to move to stand_by pose, retrying...")
                        return  # 根据具体的控制逻辑可能需要 `continue` 或其他操作

                if left_pose is not None:
                    task_state = "move_above_paper"  # 进入移动到纸巾上方的状态
                else:
                    self.get_logger().info("Waiting for target pose...")

            elif task_state == "move_above_paper":
                self.get_logger().info("Task state: move_above_paper")
                adjusted_pose = self.adjust_pose_for_paper(left_pose)
                above_pose = copy.deepcopy(adjusted_pose)
                above_pose.pose.position.z += 0.10
                if self.move_to_pose(above_pose, "right_tool_link"):
                    task_state = "move_to_paper"
                    time.sleep(1)
                else:
                    self.get_logger().warn("Failed to move above paper, retrying...")

            elif task_state == "move_to_paper":
                self.get_logger().info("Task state: move_to_paper")
                adjusted_pose = self.adjust_pose_for_paper(left_pose)
                if self.move_to_pose(adjusted_pose, "right_tool_link"):
                    task_state = "grasp"
                    time.sleep(1)
                else:
                    self.get_logger().warn("Failed to move to paper pose, retrying...")

            elif task_state == "grasp":
                self.get_logger().info("Task state: grasp")
                msg_hand_control = String()
                msg_hand_control.data = "paper_grasp"
                hand_control_publisher.publish(msg_hand_control)
                self.get_logger().info("Sent hand close command for paper.")
                time.sleep(2)
                task_state = "move_up_after_grasp"
                time.sleep(1)

            elif task_state == "move_up_after_grasp":
                self.get_logger().info("Task state: move_up_after_grasp")
                current_pose = move_group_client.get_current_pose().pose
                target_pose = copy.deepcopy(current_pose)
                target_pose.position.z += 0.10
                target_pose_stamped = PoseStamped()
                target_pose_stamped.header.frame_id = "base_link"
                target_pose_stamped.pose = target_pose
                if self.move_to_pose(target_pose_stamped, "right_tool_link"):
                    task_state = "move_to_release"
                    time.sleep(1)
                else:
                    self.get_logger().warn("Failed to move up after grasp, retrying...")

            elif task_state == "move_to_release":
                self.get_logger().info("Task state: move_to_release")
                if self.move_to_pose(release_pose, "right_tool_link"):
                    task_state = "release"
                    time.sleep(1)
                else:
                    self.get_logger().warn("Failed to move to release pose, retrying...")

            elif task_state == "release":
                self.get_logger().info("Task state: release")
                msg_hand_control = String()
                msg_hand_control.data = "release"
                hand_control_publisher.publish(msg_hand_control)
                self.get_logger().info("Sent hand open command for paper.")
                time.sleep(2)
                task_state = "return_home_after_paper"
                time.sleep(1)

            elif task_state == "return_home_after_paper":
                self.get_logger().info("Task state: return_home_after_paper")
                if self.move_to_stand_by_pose():
                    self.get_logger().info("Returned to stand_by pose after paper task.")
                    task_state = "move_above_yogurt"
                    time.sleep(1)
                else:
                    self.get_logger().warn("Failed to return to stand_by pose after paper task, retrying...")

            elif task_state == "move_above_yogurt":
                self.get_logger().info("Task state: move_above_yogurt")
                if yogurt_pose is None:
                    self.get_logger().info("Waiting for yogurt pose...")
                    rate.sleep()
                    continue
                adjusted_yogurt_pose = self.adjust_pose_for_yogurt(yogurt_pose)
                above_pose = copy.deepcopy(adjusted_yogurt_pose)
                above_pose.pose.position.x += 0.01
                above_pose.pose.position.y += -0.10
                above_pose.pose.position.z += 0.16
                if self.move_to_pose(above_pose, "right_tool_link", is_yogurt=True):
                    task_state = "move_to_yogurt"
                    time.sleep(1)
                else:
                    self.get_logger().warn("Failed to move above yogurt, retrying...")

            elif task_state == "move_to_yogurt":
                self.get_logger().info("Task state: move_to_yogurt")
                adjusted_yogurt_pose = self.adjust_pose_for_yogurt(yogurt_pose)
                if self.move_to_pose(adjusted_yogurt_pose, "right_tool_link", is_yogurt=True):
                    task_state = "grasp_yogurt"
                    time.sleep(1)
                else:
                    self.get_logger().warn("Failed to move to yogurt pose, retrying...")

            elif task_state == "grasp_yogurt":
                self.get_logger().info("Task state: grasp_yogurt")
                msg_hand_control = String()
                msg_hand_control.data = "yogurt_grasp"
                hand_control_publisher.publish(msg_hand_control)
                self.get_logger().info("Sent hand close command for yogurt.")
                time.sleep(2)
                task_state = "move_up_after_grasp_yogurt"
                time.sleep(1)


            elif task_state == "move_up_after_grasp_yogurt":
                self.get_logger().info("Task state: move_up_after_grasp_yogurt")
                current_pose = move_group_client.get_current_pose().pose
                target_pose = copy.deepcopy(current_pose)
                target_pose.position.z += 0.15
                target_pose_stamped = PoseStamped()
                target_pose_stamped.header.frame_id = "base_link"
                target_pose_stamped.pose = target_pose
                if self.move_to_pose(target_pose_stamped, "right_tool_link", is_yogurt=True):
                    task_state = "move_to_yogurt_release"
                    time.sleep(1)
                else:
                    self.get_logger().warn("Failed to move up after yogurt grasp, retrying...")

            elif task_state == "move_to_yogurt_release":
                self.get_logger().info("Task state: move_to_yogurt_release")
                if self.move_to_pose(release_pose, "right_tool_link", is_yogurt=True):
                    task_state = "release_yogurt"
                    time.sleep(1)
                else:
                    self.get_logger().warn("Failed to move to yogurt release pose, retrying...")

            elif task_state == "release_yogurt":
                self.get_logger().info("Task state: release_yogurt")
                msg_hand_control = String()
                msg_hand_control.data = "release"
                hand_control_publisher.publish(msg_hand_control)
                self.get_logger().info("Sent hand open command for yogurt.")
                time.sleep(2)
                task_state = "return_home_after_yogurt"
                time.sleep(1)

            elif task_state == "return_home_after_yogurt":
                self.get_logger().info("Task state: return_home_after_yogurt")
                if self.move_to_stand_by_pose():
                    self.get_logger().info("Task completed, returning to stand_by state.")
                    task_state = "stand_by"
                    start_move = False
                    global last_valid_pose, last_valid_pose_yogurt
                    last_valid_pose = None
                    last_valid_pose_yogurt = None
                    time.sleep(1)
                else:
                    self.get_logger().warn("Failed to return to stand_by pose after yogurt task, retrying...")

            rate.sleep()

def main(args=None):
    rclpy.init(args=args)
    try:
        task_manager_node = TaskManager()
        task_manager_node.task_manager()
    except KeyboardInterrupt:
        pass
    except ExternalShutdownException:
        pass
    finally:
        rclpy.try_shutdown()
        print("Program exited.")

if __name__ == "__main__":
    main()